import os
import json
import torch
if False: # ascend device 
    import torch_npu
    from torch_npu.contrib import transfer_to_npu   
    
import numpy as np
import torch.nn.functional as F
from PIL import Image
from torchvision import transforms
from models.hrnet import HighResolutionNet as hrnet
from tqdm import tqdm

# ---------------------------
IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)


# -------------
# 预处理函数
# -------------
def preprocess(image_path):
    """
    1) 打开图像(PIL)
    2) 转成Tensor、执行自定义增强(Gamma, Bilateral, CLAHE) + Normalize
    3) 返回增强后的Tensor
    """
    pil_image = Image.open(image_path).convert('RGB')

    # 组合所有需要的预处理/增强操作
    enhance_transform = transforms.Compose([
        transforms.ToTensor(),  # 转为 [0,1] float
        transforms.Normalize(mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD)
    ])
    
    enhanced_tensor = enhance_transform(pil_image)
    return enhanced_tensor  # C x H x W Tensor


# -------------
# 模型推理函数
# -------------
def model_prediction(model, img_tensor, mask=None,device='npu'):
    """
    1) 前向推理，返回 model 的输出(例如 segmentation map, 或者是1通道概率等)
    2) 这里可结合原始代码1:
       - sigmoid
       - 与mask相乘
       - 计算最大激活值
       - 返回预测label(二分类)
    """
    # img_tensor shape: [C, H, W]，需要给模型 batch 维度
    img_tensor = img_tensor.unsqueeze(0)  # => [1, C, H, W]

    with torch.no_grad():
        output = model(img_tensor.to(device))  # => [1, 1, H, W]
        output = torch.sigmoid(output)  # 对应原始代码中的 sigmoid
        output = output.detach().cpu() # 将输出转到CPU
        # 如果是分割输出(1通道), output shape => [1, 1, H, W]
        # 如果mask不为None，则乘上mask
        if mask is not None:
            # mask shape [1, H, W] => [1,1,H,W]
            mask = mask.unsqueeze(0) if mask.ndim == 3 else mask
            output = output * mask

        # 得到最大值
        max_val = float(torch.max(output))

        # 简单阈值判别 (原始代码1中 if max_val>0.5 => pred=1)
        prediction_bin = 1 if max_val > 0.5 else 0

    return prediction_bin, max_val, output.squeeze(0)  # 返回(标签, 最大激活值, 输出Tensor)


# -------------
# 加载模型函数
# -------------
def load_model(model_path,model_config):
    """
    参考原始代码1:
    1) 创建模型(此处仅举例，也可从config中解析)
    2) 加载权重
    3) 返回model
    """


    model = hrnet(model_config)
    # 加载预训练权重
    if model_path and os.path.exists(model_path):
        state_dict = torch.load(model_path, map_location='cpu')
        model.load_state_dict(state_dict)
        print(f"Loaded model weights from {model_path}")
    else:
        raise ValueError("Warning: model_path is None or not exist, using random init model.")

    return model


# -------------
# 主流程
# -------------
def main(data_path, split_name,model_cfg_path, model_save_path=None, save_dir=None):
    """
    1) 加载 annotations
    2) 加载 test split 列表
    3) 加载模型
    4) 模型 eval
    5) 遍历 split_list，预处理图像+生成遮罩+推理
    6) (可选)保存可视化结果到 save_dir
    7) 统计正样本数量，写入 prediction.json
    """
    # 读取 annotations
    with open(os.path.join(data_path, 'annotations.json'), 'r', encoding='utf-8') as f:
        data_dict = json.load(f)

    # 读取测试集列表
    with open(os.path.join(data_path, 'split', f"{split_name}.json"), 'r', encoding='utf-8') as f:
        split_list = json.load(f)['test']

    with open(model_cfg_path, 'r', encoding='utf-8') as f:
        model_config = json.load(f)['model']
    # 加载模型
    device='npu'
    model = load_model(model_save_path,  model_config).to(device)
    model.eval()

    print(f"model in eval mode, using device: {device}")
    # 用于统计正样本预测数量
    positive_image_cnt = 0
    prediction_label_dict = {}

    # 如果需要保存可视化mask，先确保目录存在
    if save_dir and not os.path.exists(save_dir):
        os.makedirs(save_dir, exist_ok=True)

    # 遍历测试图像
    for image_name in tqdm(split_list, desc="Processing images"):
        data_info = data_dict[image_name]
        image_path = os.path.join(data_path,data_info['enhanced_path'])

        # 1) 预处理(增强)
        enhanced_image_tensor = preprocess(image_path)  # shape [C,H,W]

        # 2) 生成遮罩(可根据你的需求决定用原图生成还是增强后生成)
        mask=Image.open(os.path.join(data_path,data_dict[image_name]['mask_path'])).resize((400,300),resample=Image.Resampling.BILINEAR)
        mask=torch.tensor(np.array(mask))
        # 3) 模型推理
        prediction_bin, max_val, prediction_map = model_prediction(model, enhanced_image_tensor, mask=mask)

        # (可选) 保存推理可视化到 save_dir
        if save_dir:
            # 这里演示将 prediction_map(1,H,W) resize到 224x224 以PNG格式保存
            # 你也可以保留原尺寸
            pred_resized = F.interpolate(prediction_map.unsqueeze(0), 
                                         size=(224, 224),
                                         mode='nearest')
            pred_resized = pred_resized.squeeze().cpu().numpy()  # (224,224)

            # 将其转为灰度图并保存
            save_path = os.path.join(save_dir, f"{image_name}.png")
            pred_img = (pred_resized * 255).astype(np.uint8)
            Image.fromarray(pred_img).save(save_path)

        # 4) 判断是否预测为正样本
        prediction_label_dict[image_name] = prediction_bin
        if prediction_bin == 1:
            positive_image_cnt += 1

    # 写出预测结果
    with open('./prediction.json', 'w', encoding='utf-8') as f:
        json.dump(prediction_label_dict, f, indent=2)

    print(f"Postive image count: {positive_image_cnt}")


# -------------
# 脚本入口
# -------------
if __name__ == "__main__":
    data_path = '/mnt/d/Workspace/MiniReti/'
    split_name = 'all'
    model_save_path = './pretrained/clr_1_hrnet_large.pth'
    save_dir = '/mnt/d/Workspace/MiniReti/ROP_ridge_segmentation'
    model_cfg_path='./config_file/hrnet_w48.json'
    main(data_path, split_name,
         model_cfg_path=model_cfg_path,
         model_save_path=model_save_path, 
         save_dir=save_dir)
